package com.rrd.hive.udf.brickhouse.udf.collect;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.MapObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector.Category;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;

/**
 * Given a Map, return an Array of structs
 * containing key and value
 */

@Description(name = "map_key_values",
        value = "_FUNC_(map) - Returns a Array of key-value pairs contained in a Map")
public class MapKeyValuesUDF extends GenericUDF {
    private MapObjectInspector moi;

    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments)
            throws UDFArgumentException {
        if (arguments.length != 1) {
            throw new UDFArgumentException("Usage : map_key_values( map) ");
        }
        if (!arguments[0].getCategory().equals(Category.MAP)) {
            throw new UDFArgumentException("Usage : map_key_values( map) ");
        }

        this.moi = (MapObjectInspector) arguments[0];

        ////
        List<String> structFieldNames = new ArrayList<>();
        List<ObjectInspector> structFieldObjectInspectors = new ArrayList<>();
        structFieldNames.add("key");
        structFieldObjectInspectors.add(this.moi.getMapKeyObjectInspector());
        structFieldNames.add("value");
        structFieldObjectInspectors.add(this.moi.getMapValueObjectInspector());

        ObjectInspector keyOI = ObjectInspectorFactory
            .getStandardStructObjectInspector(structFieldNames,
                structFieldObjectInspectors);
        ObjectInspector arrayOI = ObjectInspectorFactory
            .getStandardListObjectInspector(keyOI);

        return arrayOI;
    }

    @Override
    public Object evaluate(DeferredObject[] arguments) throws HiveException {
        Map<?, ?> map = this.moi.getMap(arguments[0].get());
        Object[] res = new Object[map.size()];
        int i = 0;
        for (Map.Entry e : map.entrySet()) {
            res[i++] = new Object[] { e.getKey(), e.getValue() };
        }
        return res;
    }

    @Override
    public String getDisplayString(String[] children) {
        return "map_key_values( " + children[0] + " )";
    }

}
